#pragma once

#include "KeeperBasic.h"
#include "KeeperLogStore.h"
#include "NuRaft.h"

namespace keeper_server {

using KeeperLogStorePtr = nuraft::ptr<KeeperLogStore>;

struct ClusterConfigWrapper {
  int server_id;
  std::string endpoint;
  NuRaftServerConfigPtr server_config;
  NuRaftClusterConfigPtr cluster_config;
};

class KeeperStateMgr : public nuraft::state_mgr {
 public:
  KeeperStateMgr(int srv_id, const std::string& endpoint,
                 std::map<int, std::string>& servers)
      : cur_log_store_(nuraft::cs_new<KeeperLogStore>()) {
    my_id_ = srv_id;
    my_endpoint_ = endpoint;
    cluster_config_wrapper_.server_config =
        nuraft::cs_new<NuRaftServerConfig>(srv_id, 0, endpoint, "", 0, srv_id);
    cluster_config_wrapper_.server_id = srv_id;
    cluster_config_wrapper_.endpoint = endpoint;

    // Initial cluster config: contains only one server (myself).
    cluster_config_wrapper_.cluster_config =
        nuraft::cs_new<NuRaftClusterConfig>();
    cluster_config_wrapper_.cluster_config->get_servers().push_back(
        cluster_config_wrapper_.server_config);
    for (auto server : servers) {
      if (server.first == srv_id) continue;
      auto peer_config =
          nuraft::cs_new<NuRaftServerConfig>(server.first, server.second);
      cluster_config_wrapper_.cluster_config->get_servers().push_back(
          peer_config);
    }
  }

  ~KeeperStateMgr() {}

  NuRaftClusterConfigPtr load_config() {
    // Just return in-memory data in this example.
    // May require reading from disk here, if it has been written to disk.
    return cluster_config_wrapper_.cluster_config;
  }

  void save_config(const NuRaftClusterConfig& config) {
    // Just keep in memory in this example.
    // Need to write to disk here, if want to make it durable.
    NuRaftBufferPtr buf = config.serialize();
    cluster_config_wrapper_.cluster_config =
        NuRaftClusterConfig::deserialize(*buf);
  }

  void save_state(const NuRaftServerState& state) {
    // Just keep in memory in this example.
    // Need to write to disk here, if want to make it durable.
    NuRaftBufferPtr buf = state.serialize();
    saved_state_ = NuRaftServerState::deserialize(*buf);
  }

  NuRaftServerStatePtr read_state() {
    // Just return in-memory data in this example.
    // May require reading from disk here, if it has been written to disk.
    return saved_state_;
  }

  nuraft::ptr<nuraft::log_store> load_log_store() { return cur_log_store_; }

  nuraft::int32 server_id() { return my_id_; }

  void system_exit(const int exit_code) {}

  NuRaftServerConfigPtr get_srv_config() const {
    return cluster_config_wrapper_.server_config;
  }

 private:
  int my_id_;
  std::string my_endpoint_;
  KeeperLogStorePtr cur_log_store_;
  NuRaftServerStatePtr saved_state_;
  ClusterConfigWrapper cluster_config_wrapper_;
};

};  // namespace keeper_server
